import socket
import threading
import json
import cv2
import numpy as np
import logging
import os

class TCPClient:
    """
    TCP 客户端类，支持文本数据和 JPEG 图像的接收，自动运行接收线程，并处理服务端断开连接等异常情况。
    """

    def __init__(self, server_ip, server_port, mode='text'):
        """
        初始化 TCP 客户端

        参数：
            server_ip (str): 服务器 IP 地址
            server_port (int): 服务器端口号
            mode (str): 数据模式，可为 'text' 或 'jpeg'
        """
        self.server_ip = server_ip
        self.server_port = server_port
        self.mode = mode  # 模式: 'text' 表示接收文本；'jpeg' 表示接收图像
        self.socket = None
        self.connected = False
        self.recv_thread = None

    def connect(self):
        """
        建立连接并启动接收线程

        返回：
            bool: 连接是否成功
        """
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)  # 创建 TCP socket
            self.socket.connect((self.server_ip, self.server_port))  # 连接服务器
            self.connected = True
            logging.info(f"[{self.mode}] 已连接到服务器 {self.server_ip}:{self.server_port}")

            # 启动接收线程
            self.recv_thread = threading.Thread(target=self._recv_loop, daemon=True)
            self.recv_thread.start()
            # logging.info(f"[{self.server_ip}:{self.server_port}] begin receiving data from server {self.mode} mode")
            return True
        except Exception as e:
            logging.error(f"[{self.mode}] 无法连接服务器: {e}")
            return False

    def close(self):
        """
        主动关闭连接
        """
        self.connected = False
        if self.socket:
            try:
                self.socket.shutdown(socket.SHUT_RDWR)  # 尝试优雅关闭
            except Exception:
                pass  # 有可能 socket 已断开
            self.socket.close()
            logging.info(f"[{self.mode}] 连接已关闭")

    def send(self, data):
        """
        向服务器发送数据

        参数：
            data (str 或 bytes): 要发送的内容
        """
        if not self.connected:
            logging.warning(f"[{self.mode}] 当前未连接，无法发送数据")
            return
        try:
            if isinstance(data, str):
                self.socket.send(data.encode("utf-8"))  # 发送字符串
                logging.debug(f"[{self.mode}] 已发送字符串: {data}")
            elif isinstance(data, bytes):
                self.socket.send(data)  # 发送二进制数据
                logging.debug(f"[{self.mode}] 已发送二进制数据 ({len(data)} 字节)")
            else:
                logging.warning(f"[{self.mode}] 不支持的发送数据类型: {type(data)}")
        except Exception as e:
            logging.error(f"[{self.mode}] 发送数据失败: {e}")
            self.close()

    def _recv_loop(self):
        """
        内部函数：接收线程主循环，根据模式选择不同的接收逻辑
        """
        try:
            if self.mode == 'text':
                self._recv_text()
            # elif self.mode == 'jpeg':
            #     self._recv_jpeg()
            else:
                logging.error(f"[{self.mode}] 无效的接收模式")
                self.close()
        except Exception as e:
            logging.error(f"[{self.mode}] 接收线程异常退出: {e}")   

    def _recv_text(self):
        """
        内部函数：接收并解析文本数据（支持 JSON）
        """
        try:
            while self.connected:
                data = self.socket.recv(1024)
                if not data:
                    logging.warning(f"[{self.mode}] 服务器关闭了连接")
                    self.close()
                    break

                message = data.decode("utf-8").strip()
                logging.debug(f"[{self.mode}] 接收到原始文本: {message}")
                try:
                    parsed = json.loads(message)  # 正确解析 JSON 对象
                    msg_type = parsed.pop("mode", None)  # 取出 type 字段

                    if msg_type == "message":
                        logging.info(f"[{self.mode}] 接收到 message 类型")
                        self.store_message(parsed)

                    elif msg_type == "control":
                        logging.info(f"[{self.mode}] 接收到 control 类型")
                        self.handle_control(parsed)

                    else:
                        logging.warning(f"[{self.mode}] 未知消息类型: {msg_type}")

                    logging.debug(f"[{self.mode}] 解析后的 JSON 数据: {parsed}")

                except json.JSONDecodeError as e:
                    logging.warning(f"[{self.mode}] JSON 解析失败: {e}")
        except (ConnectionResetError, ConnectionAbortedError) as e:
            logging.error(f"[{self.mode}] 连接被重置或中止: {e}")
            self.close()
        except Exception as e:
            logging.error(f"[{self.mode}] 接收错误: {e}")
            self.close()


    def store_message(self, new_data: dict, file_path="message_data.json"):
        """
        示例：存储消息

            new_data (dict): 要存储的消息
            file_path (str): 存储文件的路径
            message (str): 要存储的消息
        """
        # 1. 如果文件存在则读取旧数据，否则初始化为空列表
        if os.path.exists(file_path):
            with open(file_path, "r", encoding="utf-8") as f:
                try:
                    existing_data = json.load(f)
                except json.JSONDecodeError:
                    existing_data = []
        else:
            existing_data = []

        # 2. 自动计算 num 编号
        if existing_data:
            last_num = existing_data[-1].get("num", 0)
        else:
            last_num = 0
        new_data["num"] = last_num + 1

        # 3. 添加新数据并写入文件
        existing_data.append(new_data)

        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(existing_data, f, indent=2, ensure_ascii=False)

        logging.info(f"已写入 message_data.json 第 {new_data['num']} 条数据")
    
    def handle_control(self, control_data: dict):
        """
        示例：处理控制命令
        control_data (dict): 控制命令数据
        """
        logging.info(f"[{self.mode}] 接收到控制命令: {control_data}") 
        pass

    def _recv_jpeg(self):
        """
        内部函数：接收并解析 JPEG 图像数据
        """
        buffer = bytearray()
        try:
            while self.connected:
                byte = self.socket.recv(1024)
                if not byte:
                    logging.warning(f"[{self.mode}] 服务器断开图像连接")
                    self.close()
                    break

                buffer += byte  # 累加缓冲区

                while True:
                    start = buffer.find(b'\xff\xd8')  # JPEG 起始标志
                    end = buffer.find(b'\xff\xd9', start + 2)  # JPEG 结束标志

                    if start != -1 and end != -1:
                        jpeg_data = buffer[start:end + 2]
                        buffer = buffer[end + 2:]  # 移除已处理部分

                        # 解码图像
                        image = cv2.imdecode(np.frombuffer(jpeg_data, dtype=np.uint8), cv2.IMREAD_COLOR)
                        if image is not None:
                            logging.info(f"[{self.mode}] 接收到图像 ({len(jpeg_data)} 字节)")
                            # 示例：显示图像（可注释）
                            # cv2.imshow(f"{self.server_ip}:{self.server_port}", image)
                            # if cv2.waitKey(1) & 0xFF == ord('q'):
                            #     self.close()
                        else:
                            logging.warning(f"[{self.mode}] 图像解码失败")
                    else:
                        break  # 当前缓冲区中没有完整图像
        except (ConnectionResetError, ConnectionAbortedError) as e:
            logging.error(f"[{self.mode}] 图像连接异常断开: {e}")
            self.close()
        except Exception as e:
            logging.error(f"[{self.mode}] 接收图像错误: {e}")
            self.close()
